package me.flyray.bsin.server.impl;

import static dev.langchain4j.model.openai.OpenAiModelName.GPT_3_5_TURBO;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;

import java.util.List;
import java.util.Map;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import me.flyray.bsin.facade.service.VectorDocumentService;
import me.flyray.bsin.utils.RespBodyHandler;

/**
 * @author bolei
 * @date 2023/10/23
 * @desc
 */

public class VectorDocumentServiceImpl implements VectorDocumentService {

    @Autowired
    private EmbeddingModel embeddingModel;

    @Value("${vector.datasource.host}")
    private String vectorHost;
    @Value("${vector.datasource.port}")
    private Integer vectorPort;

    @Override
    public Map<String, Object> addDocument(Map<String, Object> requestMap) {

        String text = (String) requestMap.get("text");
        String collectionName = (String) requestMap.get("collectionName");

        EmbeddingStore<TextSegment> embeddingStore = MilvusEmbeddingStore.builder()
                .host(vectorHost)
                .port(vectorPort)
                .collectionName(collectionName)
                .dimension(384)
                .build();

        // 文档路径
        // Load the document that includes the information you'd like to "chat" about with the model.
        Document document = new Document(text, new Metadata());

        // Split document into segments 100 tokens each
        DocumentSplitter splitter = DocumentSplitters.recursive(
                100,
                0,
                new OpenAiTokenizer(GPT_3_5_TURBO)
        );

        List<TextSegment> segments = splitter.split(document);
        for (TextSegment segment : segments) {
            Embedding embedding1 = embeddingModel.embed(segment).content();
            embeddingStore.add(embedding1, segment);
        }

        return RespBodyHandler.RespBodyDto();
    }

    @Override
    public Map<String, Object> search(Map<String, Object> requestMap) {
        String text = (String) requestMap.get("text");
        String collectionName = (String) requestMap.get("collectionName");

        EmbeddingStore<TextSegment> embeddingStore = MilvusEmbeddingStore.builder()
                .host(vectorHost)
                .port(vectorPort)
                .collectionName(collectionName)
                .dimension(384)
                .build();

        Embedding queryEmbedding = embeddingModel.embed(text).content();
        List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.findRelevant(queryEmbedding, 1);
        EmbeddingMatch<TextSegment> embeddingMatch = relevant.get(0);
        System.out.println(embeddingMatch.score()); // 0.8144287765026093
        System.out.println(embeddingMatch.embedded().text()); // I like football.

        return RespBodyHandler.setRespBodyDto(embeddingMatch.embedded().text());
    }
}
